import numpy as np



def matmul(obj):
	a,b = obj.comeFrom[1],obj.comeFrom[2]
	a.grad += np.dot(obj.grad,b.value.T)
	b.grad += np.dot(a.value.T,obj.grad)
	return a,b

def add(obj):
	a,b = obj.comeFrom[1],obj.comeFrom[2]
	c1,c2 = np.mean(obj.grad,axis=0),np.mean(obj.grad,axis=1)

	if a.shape==obj.shape:
		a.grad += obj.grad
	elif a.shape==c1.shape:
		a.grad += c1
	elif a.shape==c2.shape:
		a.grad += c2
	else:
		a.grad +=np.mean(obj.grad)

	if b.shape==obj.shape:
		b.grad += obj.grad
	elif b.shape==c1.shape:
		b.grad += c1
	elif b.shape==c2.shape:
		b.grad += c2
	else:
		b.grad +=np.mean(obj.grad)

	return a,b

def ReLU(obj):
	a = obj.comeFrom[1]
	cache = obj.value
	cache[cache!=0] = 1
	a.grad += cache*obj.grad
	return a

def sigmoid(obj):
	a = obj.comeFrom[1]
	a.grad += obj.value*(1-obj.value)*obj.grad
	return a

def tanh(obj):
	a = obj.comeFrom[1]
	a.grad += (1-(obj.value**2))*obj.grad
	return a

def softplus(obj):
	a = obj.comeFrom[1]
	a.grad += (a.exp/(1+a.exp))*obj.grad
	return a

def simpleConv2d(obj):
	a,kernels,padding,img,stride = obj.comeFrom[1],obj.comeFrom[2],obj.comeFrom[3],obj.comeFrom[4],obj.comeFrom[5]
	kernels.grad = np.zeros(kernels.shape)
	imggrad = np.zeros(img.shape)
	out_div_in = obj.shape[1]//a.shape[1]

	for channel_in in range(a.shape[1]):
		for kernelIndex in range(out_div_in):
			channel_out = channel_in*out_div_in+kernelIndex
			kernel = kernels.value[channel_out]
			for row in range(kernel.shape[0]):
				arow = kernel[row]
				for col in range(kernel.shape[1]):
					num = arow[col]
					imggrad[:,channel_in,np.newaxis,row:row+img.shape[2]-kernel.shape[0]+1,col:col+img.shape[2]-kernel.shape[1]+1] += obj.grad[:,channel_out,np.newaxis]*kernel[row,col]
					kernels.grad[channel_out,row,col] = np.mean(img[:,channel_in,np.newaxis,row:row+img.shape[2]-kernel.shape[0]+1:stride,col:col+img.shape[2]-kernel.shape[1]+1:stride]*obj.grad[:,channel_out,np.newaxis])
					
	a.grad += imggrad[:,:,padding:padding+a.shape[2]:stride,padding:padding+a.shape[3]:stride]
	return a



def conv2d(obj):#im2col
	a,kernel,padding,stride,img,caches = obj.comeFrom[1],obj.comeFrom[2],obj.comeFrom[3],obj.comeFrom[4],obj.comeFrom[5],obj.comeFrom[6]
	# kernelGrad = np.zeros((1,kernel.size))


	img_grad = np.zeros(img.shape)
	grad = obj.grad.reshape((obj.shape[0],obj.shape[1]*obj.shape[2]))
	
	kernelGrad = np.mean(np.dot(grad,caches[1].transpose(0,2,1)),axis=(0,1)).reshape(kernel.shape)
	imgGrad = np.dot(caches[0].T,grad[:,np.newaxis,:]).transpose(1,0,2).reshape((a.shape[0],obj.shape[1],obj.shape[2],*kernel.shape))
		
	for y in range(obj.shape[1]):
		y_start = y*stride
		y_end = y_start+kernel.shape[0]
		for x in range(obj.shape[2]):
			x_start = x*stride
			x_end = x_start+kernel.shape[1]
			img_grad[:,y_start:y_end,x_start:x_end] += imgGrad[:,y,x,:,:].reshape(a.shape[0],*kernel.shape)

	
	kernel.grad += kernelGrad
	a.grad += img_grad[:,padding:a.shape[1]+padding,padding:padding+a.shape[2]]
	return a





def maxPool2d(obj):
	a,poolSize = obj.comeFrom[1],obj.comeFrom[2]
	a.grad = np.zeros(a.shape)
	out_row = 0
	out_col = 0

	for row in range(0,a.shape[2],poolSize[0]):
		for col in range(0,a.shape[3],poolSize[1]):
			a.grad[a.value[:,:,row:row+poolSize[0],col:col+poolSize[1]]==obj.grad[:,:,out_row,out_col]] += obj.grad[:,:,out_row,out_col]
		out_row += 1
		out_col = 0
	out_row = 0
	return a

def avgPool2d(obj):
	a,poolSize = obj.comeFrom[1],obj.comeFrom[2]
	a.grad = np.zeros(a.shape)
	out_row = 0
	out_col = 0

	for row in range(0,a.shape[2],poolSize[0]):
		for col in range(0,a.shape[3],poolSize[1]):
			a.grad[:,:,row:row+poolSize[0],col:col+poolSize[1]] += obj.grad[:,:,out_row,out_col]
		out_row += 1
		out_col = 0
	out_row = 0
	return a

def reshape(obj):
	a = obj.comeFrom[1]
	a.grad = self.grad.reshape(a.shape)
	return a


def softmax_crossEntropy(obj):
	a,feature = obj.comeFrom[1],obj.comeFrom[2]
	a.grad += feature.value/a.expsum*a.exp-feature.value
	return a

def MSELoss(obj):
	a,feature = obj.comeFrom[1],obj.comeFrom[2]
	a.grad += a.value-feature.value
	return a

def softsign(obj):
	a = obj.comeFrom[1]
	a.grad += (1/(1+abs(a.value))**2)*obj.grad
	return a





def f(obj):
	pass









funcDict = {'matmul':matmul,
			'add':add,
			'simpleConv2d':simpleConv2d,
			'conv2d':conv2d,
			'maxPool2d':maxPool2d,
			'avgPool2d':avgPool2d,
			'reshape':reshape,
			'ReLU':ReLU,
			'sigmoid':sigmoid,
			'tanh':tanh,
			'softplus':softplus,
			'softsign':softsign,
			'MSELoss':MSELoss,
			'softmax_crossEntropy':softmax_crossEntropy
			}







